import numpy as np
import pyworld
import glob
from hparams import hparams
import librosa
import os
def feature_world(wav,para):
    fs = para.fs
    wav = wav.astype(np.float64)
    f0, timeaxis = pyworld.harvest(wav, fs, frame_period=para.frame_period, f0_floor=71.0, f0_ceil=800.0)
    
    sp = pyworld.cheaptrick(wav, f0, timeaxis, fs)
    ap = pyworld.d4c(wav, f0, timeaxis, fs)
    coded_sp =pyworld.code_spectral_envelope(sp, fs, para.coded_dim)
    return f0,timeaxis,sp,ap,coded_sp
    
def wav_normlize(wav):
    # 信号正则
    max_ = np.max(wav)
    min_ = np.min(wav)
    wav_norm = wav*(2/(max_ - min_)) - (max_+min_)/(max_-min_)
    return wav_norm
    

def processing_wavs(file_wavs,para):
    
    f0s = []
    coded_sps = []
    for file in file_wavs:
        print("processing %s"%(file))
        # 读取音频文件
        fs = para.fs
        wav, _ = librosa.load(file, sr=fs, mono=True)
        wav = wav_normlize(wav)
        
        # 提取world 特征,采集f0和coded_sp
        f0,_,_,_,coded_sp=feature_world(wav,para)
        print(coded_sp.shape)
        f0s.append(f0)
        coded_sps.append(coded_sp)
        
    
    # 注意这里使用np.ma.log 会自动屏蔽无效值，在在计算均值时不计算
    # >>> a = np.ma.log(np.array([0,2,3]))
    # >>> a
    # masked_array(data=[--, 0.6931471805599453, 1.0986122886681098],
             # mask=[ True, False, False],
       # fill_value=1e+20)
    # >>> c = a.mean()
    # >>> c
    # 0.8958797346140275
    
    # 计算log_f0的 均值和std
    log_f0s = np.ma.log(np.concatenate(f0s))
    log_f0s_mean = log_f0s.mean()
    log_f0s_std = log_f0s.std()
    
    # 计算 coded_sp 的均值和 标准差
    coded_sps_array = np.concatenate(coded_sps,axis=0)  # coded_sp的维度  T * D
    coded_sps_mean = np.mean(coded_sps_array,axis=0,keepdims = True)
    coded_sps_std = np.std(coded_sps_array,axis=0,keepdims = True)

    # 利用 coded_sp 的均值和 标准差 对特征进行正则
    coded_sps_norm = []
    for coded_sp in coded_sps:
        coded_sps_norm.append(  (coded_sp- coded_sps_mean)/ coded_sps_std )
        
    return log_f0s_mean,log_f0s_std,coded_sps_mean,coded_sps_std,coded_sps_norm

if __name__ == "__main__":
    
    para = hparams()
    
    # 提取说话人A的特征
    dir_train_A = para.train_dir_A
    wavs = glob.glob(dir_train_A+'/*wav') 
    f0_mean,f0_std,mecp_mean,mecp_std, mecps = processing_wavs(wavs,para)
   
    os.makedirs(para.catch_A,exist_ok = True)
    np.save(os.path.join(para.catch_A,'static_f0.npy'),np.array([f0_mean,f0_std],dtype=object))
    np.save(os.path.join(para.catch_A,'static_mecp.npy'),np.array([mecp_mean,mecp_std],dtype=object))
    np.save(os.path.join(para.catch_A,'data.npy'),np.array(mecps,dtype=object))
    
    
    # 提取说话人B的特征
    dir_train_B = para.train_dir_B
    wavs = glob.glob(dir_train_B+'/*wav')
    f0_mean,f0_std,mecp_mean,mecp_std, mecps = processing_wavs(wavs,para)
    os.makedirs(para.catch_B,exist_ok = True)
    np.save(os.path.join(para.catch_B,'static_f0.npy'),np.array([f0_mean,f0_std],dtype=object))
    np.save(os.path.join(para.catch_B,'static_mecp.npy'),np.array([mecp_mean,mecp_std],dtype = object))
    np.save(os.path.join(para.catch_B,'data.npy'),np.array(mecps,dtype=object))
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    
    